from typing import TypedDict
from langgraph.graph import START, StateGraph
from langchain_community.chat_models import ChatTongyi

import os
os.environ["DASHSCOPE_API_KEY"] = os.getenv("OPENAI_API_KEY")

model = ChatTongyi(model="qwen-plus")


class State(TypedDict):
      topic: str
      joke: str
      poem: str


def write_joke(state: State):
      topic = state["topic"]
      joke_response = model.invoke(
            [{"role": "user", "content": f"Write a joke about {topic}"}]
      )
      return {"joke": joke_response.content}


def write_poem(state: State):
      topic = state["topic"]
      poem_response = model.invoke(
            [{"role": "user", "content": f"Write a short poem about {topic}"}]
      )
      return {"poem": poem_response.content}


graph = (
      StateGraph(State)
      .add_node(write_joke)
      .add_node(write_poem)
      # 并行启动
      .add_edge(START, "write_joke")
      .add_edge(START, "write_poem")
      .compile()
)

for msg, metadata in graph.stream(
    {"topic": "cats"},
    stream_mode="messages",
):
    if msg.content and metadata["langgraph_node"] == "write_poem":
        print(msg.content, end="|", flush=True)